# -*- coding: utf-8 -*-
# @Time     : 2020/7/23 1:56 下午
# @Author   : huangxiong
# @FileName : stress_testing.py
# @Comment  :
# @Software : PyCharm
import pandas as pd
import numpy as np
import statsmodels.api as sm

from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.datasource_fetch.portfolio_api.portfolio_tool import get_portfolio_fund_weight
from quant_researcher.quant.risk_management.fof_related.VaR_analysis import FUND_NAME_TABLE
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool.time_tool import date_shifter, get_yesterday
from quant_researcher.quant.project_tool.exception import FOFServerError, ERR_S_2

FUND_RET_TABLE = 'mf_di_fnddailyreturn'
FUND_NAV_TABLE = 'mf_di_fndadjnav'
FUND_BENCHMARK_TABLE = 'mf_di_fndbenchmark'
INDEX_RET_TABLE = 'mf_di_idxdailyreturn'


def reg_beta(fund_list, index_code, fund_where_condition, fund_weight_df):
    conn = db_conn.get_derivative_data_conn()

    # 计算贝塔时，每次都使用代码运行时间的前一天作为计算结束日，计算结束日一年前作为计算开始日
    calc_date_end = get_yesterday()
    calc_date_start = date_shifter(calc_date_end, 'years', -1)

    # 取近一年的基金收益率序列
    fund_ret_4_reg = pd.read_sql(f"select fund_code, day as tj, daily_return as ret "
                                 f"from {FUND_RET_TABLE} "
                                 f"where {fund_where_condition} "
                                 f"and day >= '{calc_date_start}' "
                                 f"and day <= '{calc_date_end}' ", conn)
    fund_ret_4_reg['tj'] = fund_ret_4_reg['tj'].astype(str)
    fund_ret_4_reg = fund_ret_4_reg.set_index(['tj', 'fund_code'])['ret'].unstack()

    # 取各基金对应的基准以及基准的收益率序列
    fund_bench = pd.read_sql(f"select fund_code, benchmark "
                             f"from {FUND_BENCHMARK_TABLE} "
                             f"where {fund_where_condition} "
                             f"and end_date = (select max(end_date) from mf_di_fndbenchmark) ", conn)
    if len(fund_list) != len(list(fund_bench['fund_code'])):
        LOG.error(f"组合包含的基金为 {fund_list} ，有对应基准的基金为 {list(fund_bench['fund_code'])}")
        conn.close()
        return None, None
    bench_list = fund_bench['benchmark'].tolist()
    bench_where_condition = get_lst_condition(bench_list, 'index_code')
    bench_ret = pd.read_sql(f"select index_code as benchmark, day as tj, daily_return as ret "
                            f"from {INDEX_RET_TABLE} "
                            f"where {bench_where_condition} "
                            f"and day >= '{calc_date_start}' "
                            f"and day <= '{calc_date_end}' ", conn)
    bench_ret['tj'] = bench_ret['tj'].astype(str)
    bench_ret = bench_ret.set_index(['tj', 'benchmark'])['ret'].unstack()
    # 构造与基金对应的基准收益率DataFrame
    bench_ret_df = pd.DataFrame(index=bench_ret.index, columns=fund_list)
    for fund in fund_list:
        bench_ret_df[fund] = bench_ret[fund_bench[fund_bench['fund_code'] == fund]['benchmark']]

    # 如果某支基金的收益率序列小于10个数据，则将该基金收益率序列替换为其对应基准的收益率序列
    fund_ret_count = fund_ret_4_reg.count()
    fund_2_replace = list(fund_ret_count[fund_ret_count < 10].index)
    if len(fund_2_replace) > 0:
        fund_ret_4_reg[fund_2_replace] = np.nan
        fund_ret_4_reg[fund_2_replace] = fund_ret_4_reg[fund_2_replace].fillna(bench_ret_df)
    fund_ret_4_reg = fund_ret_4_reg.dropna()
    if fund_ret_4_reg.shape[0] < 10:
        LOG.error(f"dropna之后的基金收益率序列个数小于10个，应该是出现了数据问题")
        conn.close()
        return None, None

    # 找到index_code的收益率序列
    index_ret_4_reg = pd.read_sql(f"select day as tj, daily_return as ret "
                                  f"from {INDEX_RET_TABLE} "
                                  f"where index_code = '{index_code}' "
                                  f"and day >= '{calc_date_start}' "
                                  f"and day <= '{calc_date_end}' ", conn)
    index_ret_4_reg['tj'] = index_ret_4_reg['tj'].astype(str)
    index_ret_4_reg = index_ret_4_reg.set_index('tj').rename(columns={'ret': 'index'})
    all_ret_4_reg = fund_ret_4_reg.merge(index_ret_4_reg, how='left', on='tj')
    all_ret_4_reg = all_ret_4_reg.fillna(0)
    reg_count_row = all_ret_4_reg.shape[0]

    # 回归计算beta
    X = sm.add_constant(all_ret_4_reg['index'])
    Y = all_ret_4_reg[all_ret_4_reg.columns.difference(['index'])]
    result = sm.OLS(Y, X, missing='drop').fit()
    if Y.shape[1] > 1:
        beta_alpha = result.params.T
    else:
        beta_alpha = result.params.to_frame().T
    beta_alpha.index = Y.columns
    beta_alpha = beta_alpha.rename(columns={'const': 'alpha', 'index': 'm_beta'})

    fund_beta = beta_alpha.reset_index().rename(columns={'index': 'fund_code'})
    fund_beta = fund_beta.merge(fund_weight_df, how='left', on='fund_code')
    fund_beta['beta'] = fund_beta['m_beta'] * fund_beta['weight']

    conn.close()

    return fund_beta, reg_count_row


def pressure_test_history(portfolio_id, index_code, start_date, end_date, calc_date):
    """
    压测情景为历史时段

    :param portfolio_id: str，组合id
    :param index_code: str，基准指数代码
    :param start_date: str，开始日期，格式形如'2019-07-09'
    :param end_date: str，结束日期，格式形如'2020-07-09'
    :param calc_date: str，计算日期，格式形如'2020-07-21'，目前该参数没有用到
    :return:
    """
    # 根据组合代码找到持有的基金及权重
    tmp_df = get_portfolio_fund_weight(portfolio_id, end_date.replace('-', ''))
    if tmp_df is None:
        raise FOFServerError(ERR_S_2, f"该组合 {portfolio_id} 的id有问题")
    fund_list = tmp_df['fund_code'].tolist()
    fund_weight_df = tmp_df.rename(columns={'fund_weight': 'weight'})
    fund_where_condition = get_lst_condition(fund_list, 'fund_code')

    fund_beta, reg_count_row = reg_beta(fund_list, index_code, fund_where_condition, fund_weight_df)
    if fund_beta is None:
        raise FOFServerError(ERR_S_2, f"基金{fund_list}回归计算beta时出现问题")

    conn_base = db_conn.get_basic_data_conn()
    conn_ty = db_conn.get_derivative_data_conn()
    # 计算start_date到end_date时间段内基金的收益率和基准的收益率
    fund_nav = pd.read_sql(f"select fund_code, nav_date as tj, adjusted_nav as nav "
                           f"from {FUND_NAV_TABLE} "
                           f"where {fund_where_condition} "
                           f"and nav_date >= '{start_date}' "
                           f"and nav_date <= '{end_date}' "
                           f"order by nav_date ", conn_ty)
    fund_nav['tj'] = fund_nav['tj'].astype(str)
    fund_nav = fund_nav.set_index(['tj', 'fund_code'])['nav'].unstack()
    fund_nav = fund_nav.ffill().bfill()
    fund_ret = (fund_nav.iloc[-1] / fund_nav.iloc[0] - 1).rename('m_ret').reset_index()

    index_price = pd.read_sql(f"select day as tj, price_to as price "
                              f"from {INDEX_RET_TABLE} "
                              f"where index_code = '{index_code}' "
                              f"and day >= '{start_date}' " 
                              f"and day <= '{end_date}' "
                              f"order by day ", conn_ty)
    index_ret = index_price['price'].iloc[-1] / index_price['price'].iloc[0] - 1

    # 计算基金的收益边际贡献和收益
    fund_ret = fund_ret.merge(fund_weight_df, how='left', on='fund_code')
    fund_ret['ret'] = fund_ret['m_ret'] * fund_ret['weight']
    fund_ret['m_a_ret'] = fund_ret['m_ret'] - index_ret
    fund_ret['a_ret'] = fund_ret['m_a_ret'] * fund_ret['weight']
    fund_ret = fund_ret.set_index('fund_code')
    fund_ret = fund_ret.round(6)
    fund_ret = fund_ret.reset_index()

    # 计算压力测试相关参数
    # 压力测试边际贡献
    fund_beta['p_t_m_c'] = fund_beta['m_beta'] * index_ret + fund_beta['alpha'] * reg_count_row
    # 压力测试贡献
    fund_beta['p_t_c'] = fund_beta['p_t_m_c'] * fund_beta['weight']
    # 压力测试主动边际贡献
    fund_beta['p_t_a_m_c'] = fund_beta['p_t_m_c'] - index_ret
    # 压力测试主动贡献
    fund_beta['p_t_a_c'] = fund_beta['p_t_a_m_c'] * fund_beta['weight']
    fund_beta = fund_beta.set_index('fund_code')
    fund_beta = fund_beta.round(6)
    fund_beta = fund_beta.reset_index()

    fund_ret = fund_ret.drop(columns='weight')
    fund_beta = fund_beta.drop(columns='alpha')

    # 获取基金名称
    fund_name = pd.read_sql(f"select fund_code, fund_sname "
                            f"from {FUND_NAME_TABLE} "
                            f"where {fund_where_condition} ", conn_base)
    fund_ret = fund_ret.merge(fund_name, how='left', on='fund_code')
    fund_beta = fund_beta.merge(fund_name, how='left', on='fund_code')
    fund_ret = fund_ret.to_dict(orient='records')
    fund_beta = fund_beta.to_dict(orient='records')
    res_list = [
        {'fund_ret': fund_ret},
        {'fund_beta': fund_beta},
    ]

    conn_base.close()
    conn_ty.close()

    return res_list


def pressure_test_impact(portfolio_id, index_code, index_impact, forecast_period, calc_date):
    """
    压测情景为冲击量

    :param portfolio_id: str，组合id
    :param index_code: str，基准指数代码
    :param index_impact: str，基准冲击量，格式形如'1'
    :param forecast_period: str，预测天数，格式形如'1'
    :param calc_date: str，计算日期，格式形如'2020-07-21'，目前该参数没有用到
    :return:
    """
    # 根据组合代码找到持有的基金及权重
    tmp_df = get_portfolio_fund_weight(portfolio_id, calc_date.replace('-', ''))
    if tmp_df is None:
        raise FOFServerError(ERR_S_2, f"该组合 {portfolio_id} 的id有问题")
    fund_list = tmp_df['fund_code'].tolist()
    fund_weight_df = tmp_df.rename(columns={'fund_weight': 'weight'})
    fund_where_condition = get_lst_condition(fund_list, 'fund_code')

    fund_beta, reg_count_row = reg_beta(fund_list, index_code, fund_where_condition, fund_weight_df)

    if fund_beta is None:
        raise FOFServerError(ERR_S_2, f"基金{fund_list}回归计算beta时出现问题")

    index_impact = int(index_impact) / 100
    forecast_period = int(forecast_period)

    conn_base = db_conn.get_basic_data_conn()

    # 计算压力测试相关参数
    # 压力测试边际贡献
    fund_beta['p_t_m_c'] = fund_beta['m_beta'] * index_impact + fund_beta['alpha'] * forecast_period
    # 压力测试贡献
    fund_beta['p_t_c'] = fund_beta['p_t_m_c'] * fund_beta['weight']
    # 压力测试主动边际贡献
    fund_beta['p_t_a_m_c'] = fund_beta['p_t_m_c'] - index_impact
    # 压力测试主动贡献
    fund_beta['p_t_a_c'] = fund_beta['p_t_a_m_c'] * fund_beta['weight']

    fund_ret = fund_beta[['fund_code', 'p_t_m_c', 'p_t_c', 'p_t_a_m_c', 'p_t_a_c']]
    fund_ret = fund_ret.rename(columns={'p_t_m_c': 'm_ret', 'p_t_c': 'ret',
                                        'p_t_a_m_c': 'm_a_ret', 'p_t_a_c': 'a_ret'})
    fund_ret = fund_ret.set_index('fund_code')
    fund_ret = fund_ret.round(6)
    fund_ret = fund_ret.reset_index()

    fund_beta = fund_beta.drop(columns='alpha')
    fund_beta = fund_beta.set_index('fund_code')
    fund_beta = fund_beta.round(6)
    fund_beta = fund_beta.reset_index()

    # 获取基金名称
    fund_name = pd.read_sql(f"select fund_code, fund_sname "
                            f"from {FUND_NAME_TABLE} "
                            f"where {fund_where_condition} ", conn_base)
    fund_ret = fund_ret.merge(fund_name, how='left', on='fund_code')
    fund_beta = fund_beta.merge(fund_name, how='left', on='fund_code')

    fund_ret = fund_ret.to_dict(orient='records')
    fund_beta = fund_beta.to_dict(orient='records')
    res_list = [
        {'fund_ret': fund_ret},
        {'fund_beta': fund_beta},
    ]

    conn_base.close()

    return res_list


if __name__ == "__main__":
    # pressure_test_history(portfolio_id='291', index_code='000300',
    #                       start_date='2019-07-09', end_date='2020-07-09', calc_date='2020-07-21')
    pressure_test_impact(portfolio_id='1199', index_code='000300',
                         index_impact='3', forecast_period='3', calc_date='2020-08-27')
